package com.zyh.flink.day04;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.*;
import java.util.stream.Collectors;

public class CouponBroadcastStateJob2 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        SingleOutputStreamOperator<Tuple2<Coupon, String>> couponStream = environment.socketTextStream("hadoop10", 9991)
                .map(new MapFunction<String, Tuple2<Coupon, String>>() {
                    @Override
                    public Tuple2<Coupon, String> map(String s) throws Exception {
                        String[] ss = s.split("\\s+");
                        return Tuple2.of(new Coupon(ss[0], Double.valueOf(ss[1]), Double.valueOf(ss[2])), ss[3]);
                    }
                });

        SingleOutputStreamOperator<Order> orderStream = environment.socketTextStream("hadoop10", 9992)
                .map(new MapFunction<String, Order>() {
                    @Override
                    public Order map(String s) throws Exception {
                        String[] ss = s.split("\\s+");
                        return new Order(ss[0], ss[1], Double.valueOf(ss[2]));
                    }
                });

        KeyedStream<Order, Tuple2<String, String>> keyedStream = orderStream.keyBy(order -> Tuple2.of(order.getUserNo(), order.getCategory()),Types.TUPLE(Types.STRING,Types.STRING));

        //广播低速流
        /*
        * 广播状态key：类别 value：Set<Coupon>
        * */
        MapStateDescriptor<String, Set<Coupon>> msd = new MapStateDescriptor<>("coupons", String.class, (Class<Set<Coupon>>) (Class<?>) Set.class);
        BroadcastStream<Tuple2<Coupon, String>> broadcastStream = couponStream.broadcast(msd);

        //高速流连接广播流形成广播连接流
        BroadcastConnectedStream<Order, Tuple2<Coupon, String>> connectedStream = keyedStream.connect(broadcastStream);
        SingleOutputStreamOperator<String> result = connectedStream.process(new MyKeyedBroadcastProcessFunction2(msd));
        result.print();
        environment.execute("KeyedStream");
    }
}

/*
* KS:高速流上的分组依据类型 Tuple2<String,String>
* IN1:高速流上的元素类型Order
* IN2:广播流的数据类型；低吞吐量流的数据类型
* OUT:输出数据类型
* */
class MyKeyedBroadcastProcessFunction2 extends KeyedBroadcastProcessFunction<Tuple2<String,String>,Order,Tuple2<Coupon,String>,String>{
    private MapStateDescriptor<String,Set<Coupon>> msd;

    MyKeyedBroadcastProcessFunction2(MapStateDescriptor<String,Set<Coupon>> msd){
        this.msd = msd;
    }
    private ReducingState<Double> reducingState;

    @Override
    public void open(Configuration parameters) throws Exception {
        RuntimeContext ctx = getRuntimeContext();
        ReducingStateDescriptor<Double> rsd = new ReducingStateDescriptor<>("total", (v1, v2) -> v1 + v2, Types.DOUBLE);
        this.reducingState = ctx.getReducingState(rsd);
    }

    @Override
    public void processElement(Order order, ReadOnlyContext readOnlyContext, Collector<String> collector) throws Exception {
        this.reducingState.add(order.getMoney());
        Double total = this.reducingState.get();
        //从广播状态上获取优惠券集合
        ReadOnlyBroadcastState<String, Set<Coupon>> broadcastState = readOnlyContext.getBroadcastState(this.msd);
        Set<Coupon> couponSet = broadcastState.get(order.getCategory());

        if (couponSet!=null){
            //自高向低判断
            List<Coupon> couponList = couponSet.stream().collect(Collectors.toList());
            Collections.sort(couponList, (o1,o2)-> (int) (o2.getFull()-o1.getFull()));

            for (Coupon coupon : couponList) {
                if (total>=coupon.getFull()){
                    collector.collect(order.getUserNo()+"在类别："+order.getMoney()+"已经消费达标，满"+coupon.getFull()+"减"+coupon.getRebate());
                    this.reducingState.clear();
                    this.reducingState.add(total-coupon.getFull());
                    return;
                }
            }
            //循环结束没有满足条件的优惠券
            Coupon coupon = couponList.get(couponList.size() - 1);
            collector.collect(order.getUserNo()+"在类别"+order.getCategory()+"还差"+(coupon.getFull()-total));
        }else {
            collector.collect(order.getCategory()+"类别下没有优惠券");
        }
    }

    @Override
    public void processBroadcastElement(Tuple2<Coupon, String> value, Context context, Collector<String> collector) throws Exception {
        Coupon coupon = value.f0;
        String operator = value.f1;

        BroadcastState<String, Set<Coupon>> broadcastState = context.getBroadcastState(this.msd);

        Set<Coupon> coupons = broadcastState.get(coupon.getCategory());
        if (coupons==null){
            coupons = new HashSet<>();
        }
        //根据操作符更新优惠券信息
        if ("+".equals(operator)){
            coupons.add(coupon);
        }else if ("-".equals(operator)){
            coupons.remove(coupon);
        }
        //将更新后的集合同步到状态中
        broadcastState.put(coupon.getCategory(),coupons);
    }
}